package cn.sw.system.ai.controller;

import cn.sw.system.ai.manager.AiManager;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ModelData;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import io.reactivex.schedulers.Schedulers;
import io.reactivex.Scheduler;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

@Slf4j
@RestController
@RequestMapping("/ai")
public class AIController {

    @Resource
    private AiManager aiManager;

    private static final String AI_PROMPT =
            "你是小布，擅长解决小动物问题...\n" +
                    "<!--技能模块-->\n" +
                    "{具体提示词内容}";
    private static final List<ChatMessage> AI_INTERACTION_MESSAGE = Collections.singletonList(
            new ChatMessage("system", AI_PROMPT)
    );

    @GetMapping("/ai_generate/sse")
    public SseEmitter aiGenerateQuestionSSE(@RequestParam String question,
                                            HttpServletRequest request,
                                            HttpServletResponse response) throws Exception {

        log.info("收到 SSE 请求，问题内容: {}", question);

        if (question == null || question.trim().isEmpty()) {
            return null;
        }
//  1.设置响应头：告知浏览器这是一个事件流（SSE），编码格式为 UTF-8。
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");

//  2.构建对话历史  将系统提示词和用户问题构造成完整的对话历史，准备发送给 AI 模型。
        List<ChatMessage> messages = new ArrayList<>(AI_INTERACTION_MESSAGE);
        messages.add(new ChatMessage("user", question));

//  3.创建 SseEmitter 实例
        SseEmitter sseEmitter = new SseEmitter(60_000L);

//  4.准备响应式流处理 使用 RxJava 的 Schedulers.single() 调度器保证顺序执行。
//  5.stringBuilder 用于拼接 AI 返回的内容片段。
        Scheduler scheduler = Schedulers.single();
        StringBuilder stringBuilder = new StringBuilder();

//  6.调用 AiManager 的 doStreamRequest 方法发起流式请求，获取 AI 返回的数据流。
        Flowable<ModelData> modelDataFlowable = aiManager.doStreamRequest(messages, AiManager.STABLE_TEMPERATURE);

//  7.处理数据流
        modelDataFlowable
                .doOnSubscribe(d -> log.info("开始流式响应"))//订阅流时记录日志。
                .observeOn(scheduler)//指定在单线程调度器中处理数据。
                .map(modelData -> modelData.getChoices().get(0).getDelta().getContent())//提取 AI 返回内容中的文本片段。
                .doOnNext(content -> {
                    stringBuilder.append(content);//将内容追加到 stringBuilder 中；
                    sseEmitter.send(SseEmitter.event().name("message").data(content));//通过 sseEmitter 向前端发送事件名为 message 的数据；
                    stringBuilder.setLength(0);//清空 stringBuilder。
                })
                .doOnError((e) -> log.error("sse error", e))//发生错误时记录日志。
                .doOnComplete(sseEmitter::complete)//流完成时通知前端连接关闭。
                .subscribe();//触发整个流式处理过程。

        return sseEmitter;//返回结果
    }


}
